# 97. K倍区间
# https://www.lanqiao.cn/problems/97/learning/?page=1&first_category_id=1&contain_answer=true&second_category_id=3
# Date:2025/1/4
from leetcode.test import test_function as tf
import numpy as np


def k_fold_interval(k, nums):
    """暴力法"""
    res = 0
    for start, num in enumerate(nums):
        sum_num = num
        if sum_num % k == 0:
            res += 1
        for i in range(start + 1, len(nums)):
            sum_num += nums[i]
            if sum_num < k:
                continue
            if sum_num % k == 0:
                res += 1
    return res


def k_fold_interval_opt(k, nums):
    """使用同余定理
    如: [1 2 3 4 5] k=2
    求nums的前缀和为 [1 3 6 10 15]
    对2取余数[1 1 0 0 1]会发现余数为0的部分前缀和能被k整除, 余数相减为0的也一定是2的倍数.
    实现方法为: 前缀和 + 哈希表"""
    remainder_count = {0: 1}  # 初始化，余数为 0 的频次为 1（表示前缀和本身能被 k 整除）
    prefix_sum = 0
    res = 0

    for num in nums:
        prefix_sum += num
        remainder = prefix_sum % k
        if remainder < 0:  # 处理负数余数，转化为正数
            remainder += k
        if remainder in remainder_count:  # 如果有余数相同的情况
            res += remainder_count[remainder]
        remainder_count[remainder] = remainder_count.get(remainder, 0) + 1

    return res


def k_fold_interval_sr(k, nums):
    """使用同余定理
    如: [1 2 3 4 5] k=2
    求nums的前缀和为 [1 3 6 10 15]
    对2取余数[1 1 0 0 1]会发现余数为0的部分前缀和能被k整除, 余数相减为0的也一定是2的倍数, 两两组合即C(3,2) = 3
    """
    prefix_sum = 0
    remainder_count = [0] * k
    for num in nums:
        prefix_sum += num
        remainder_count[prefix_sum % k] += 1  # 统计前缀和额的不同余数的个数

    res = remainder_count[0]  # 统计可以被k整除的前缀和的数量
    for i in range(k):  # 计算卡特兰数
        temp = remainder_count[i]
        res += temp * (temp - 1) // 2
    return res


if __name__ == '__main__':
    # inpstr = input()
    # n, k = (int(i) for i in inpstr.split())
    # nums = []
    # for i in range(n):
    #     nums.append(int(input()))
    # print(k_fold_interval_opt(k, nums))
    np.random.seed(33)
    inp = [{"k": 2, "nums": [1, 2, 3, 4, 5]},
           {"k": 7, "nums": [1, 3, 4, 7, 9, 13]},
           {"k": 7, "nums": np.random.randint(0, high=2000, size=1000)},
           {"k": 7, "nums": np.random.randint(0, high=20000, size=10000)}, ]
    out = [6, 3, 71233, 7147248]
    tf(k_fold_interval, inp, out)  # 0s, 0.08s, 7.87s
    tf(k_fold_interval_opt, inp, out)  # 0s, 0.08s, 7.87s
    tf(k_fold_interval_sr, inp, out)  # 0s, 0.08s, 7.87s
